package com.ljh.myrpc.registry;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.cron.CronUtil;
import cn.hutool.cron.task.Task;
import cn.hutool.json.JSONUtil;
import com.ljh.myrpc.config.RegistryConfig;
import com.ljh.myrpc.model.ServiceMetaInfo;
import io.etcd.jetcd.*;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;
import lombok.extern.slf4j.Slf4j;
import org.checkerframework.checker.units.qual.C;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

@Slf4j
public class EtcdRegistry implements Registry{

    private Client client;
    private KV kvClient;
    private final Set<String> localRegisterNodeKeySet = new HashSet<>();
    // 本地服务信息缓存
    private final RegistryServiceCache registryServiceCache = new RegistryServiceCache();
    // 正在监听的key集合
    private final Set<String> watchingKeySet = new ConcurrentHashSet<>();
    /**
     * 根节点
     * ETCD键存储的根路径
     */
    private static final String ETCD_ROOT_PATH = "/rpc/";

    @Override
    public void init(RegistryConfig registryConfig) {
        client = Client.builder()
                .endpoints(registryConfig.getAddress())
                .connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
                .build();
        kvClient = client.getKVClient();
        heartBeat();
    }

    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws Exception {
        // 创建lease客户端
        Lease lease = client.getLeaseClient();
        // 创建一个30s的租约
        long id = lease.grant(30).get().getID();
        // 设置要存储的键值对
        String registryKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence key = ByteSequence.from(registryKey, StandardCharsets.UTF_8);
        ByteSequence value = ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);
        // 将键值对和租约关联，设置过期时间
        PutOption putOption = PutOption.builder().withLeaseId(id).build();
        kvClient.put(key, value, putOption).get();
        // 将注册的节点key存储到本地集合中，便于后续注销操作
        localRegisterNodeKeySet.add(registryKey);
        log.info("服务注册成功 {}-{}", registryKey, JSONUtil.toJsonStr(serviceMetaInfo));
    }

    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) {
        String registerKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence key = ByteSequence.from(
                registerKey,
                StandardCharsets.UTF_8);
        kvClient.delete(key);
        localRegisterNodeKeySet.remove(registerKey);
        log.info(key + " 服务注销成功");
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
        // 优先使用本地缓存
        List<ServiceMetaInfo> serviceMetaInfos = registryServiceCache.readCache(serviceKey);
        if(serviceMetaInfos != null && !serviceMetaInfos.isEmpty()){
            log.info("从本地缓存中获取服务信息，key: {}", serviceKey);
            return serviceMetaInfos;
        }
        // 前缀搜索
        String searchPrefix = ETCD_ROOT_PATH + serviceKey + "/";

        GetOption getOption = GetOption.builder().isPrefix(true).build();
        try {
            // 前缀查询
            List<KeyValue> kvs = kvClient.get(ByteSequence.from(searchPrefix, StandardCharsets.UTF_8), getOption)
                    .get().getKvs();
            // 将获取到的键值对转换为ServiceMetaInfo对象列表
            List<ServiceMetaInfo> serviceMetaInfoList = kvs.stream().map(kv -> {
                String k = kv.getKey().toString(StandardCharsets.UTF_8);
                watch(k); // 监听服务节点变化
                String v = kv.getValue().toString(StandardCharsets.UTF_8);
                return JSONUtil.toBean(v, ServiceMetaInfo.class);
            }).toList();
            // 写入缓存
            registryServiceCache.writeCache(serviceKey, serviceMetaInfoList);
            log.info("注册中心服务发现成功，key: {}, 服务信息: {}", serviceKey, JSONUtil.toJsonStr(serviceMetaInfoList));
            return serviceMetaInfoList;
        }catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void destory() {
        log.info("当前节点下线");
        // 遍历本节点所有的key
        for (String key : localRegisterNodeKeySet) {
            try {
                kvClient.delete(ByteSequence.from(key, StandardCharsets.UTF_8)).get();
            } catch (Exception e) {
                throw new RuntimeException(key + "节点下线失败" + e);
            }
        }
        // 释放资源
        if(kvClient != null){
            kvClient.close();
        }
        if (client != null) {
            client.close();
        }
    }

    @Override
    public void heartBeat() {
        // 10秒一次心跳
        CronUtil.schedule("*/10 * * * * *", new Task(){
            @Override
            public void execute() {
                // 遍历本节点所有的key
                for (String key : localRegisterNodeKeySet) {
                    try {
                        List<KeyValue> kvs = kvClient.get(ByteSequence.from(key, StandardCharsets.UTF_8)).get().getKvs();
                        if(CollUtil.isEmpty(kvs)){
                            // 节点已过期，需要重启节点才能重新注册
                            continue;
                        }
                        KeyValue keyValue = kvs.get(0);
                        String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                        ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
                        register(serviceMetaInfo);
                    } catch (Exception e) {
                        log.error("{} 续签失败，{}",key,e);
                        throw new RuntimeException(e);
                    }
                }
            }
        });
        // 启动定时任务
        CronUtil.setMatchSecond(true); // 支持秒级别定时任务
        CronUtil.start();
    }

    /**
     * 监听服务节点变化（消费端）
     * @param serviceNodeKey
     */
    @Override
    public void watch(String serviceNodeKey) {
        Watch watchClient = client.getWatchClient();
        //
        boolean added = watchingKeySet.add(serviceNodeKey);
        if(added){
            watchClient.watch(ByteSequence.from(serviceNodeKey, StandardCharsets.UTF_8), watchResponse -> {
                for (WatchEvent event : watchResponse.getEvents()) {
                    switch (event.getEventType()){
                        case DELETE :
                            registryServiceCache.clearCache(serviceNodeKey);
                            break;
                        case PUT:
                        default:
                            break;
                    }
                }
            });
        }
    }
}
